import os
import torch

from cross_exp.exp_crossformer import Exp_crossformer
from utils.tools import string_split

# xwei

def argDefin():
    # 直接定义参数（完全替代命令行参数）
    class Args:
        def __init__(self):
            # 核心参数（对应命令行参数），注释及取值范围参考992_crossformer_transformer_utiliz.pdf实验设置与模型设计
            self.data = 'ETTh1'  # --data：指定使用的数据集，此处为电力变压器温度小时级数据集（ETTh1）；取值范围：文档中6个真实数据集，即'ETTh1'/'ETTm1'/'WTH'/'ECL'/'ILI'/'Traffic' {insert\_element\_0\_}
            self.in_len = 168  # --in_len：输入序列长度（历史时间步数量T），表示用过去168个时间步数据预测；取值范围：ETTh1/WTH/ECL/Traffic取{24,48,96,168,336,720}，ETTm1取{24,48,96,192,288,672}，ILI取{24,36,48,60} {insert\_element\_1\_}
            self.out_len = 24  # --out_len：输出序列长度（预测时间步数量τ），表示预测未来24个时间步结果；取值范围：常见{24,48,96,168,336,720}（ETTh1/WTH/ECL/Traffic）、{24,48,96,288,672}（ETTm1）、{24,36,48,60}（ILI） {insert\_element\_2\_}
            self.seg_len = 6  # --seg_len：DSW嵌入的分段长度（L_seg），将每个维度序列按6个时间步划分为一个分段；取值范围：{4,6,12,24,48}，短期预测用小值，长期预测用大值（如ETTh1长期预测最优为24） {insert\_element\_3\_}
            self.learning_rate = 1e-4  # --learning_rate：初始学习率，控制参数更新幅度；取值范围：{5e-3,1e-3,5e-4,1e-4,5e-5,1e-5}，通过网格搜索选择 {insert\_element\_4\_}
            self.itr = 3  # --itr：实验重复运行次数，减少随机初始化影响，取多次结果平均值；取值范围：{3,5,10}，次数越多结果越稳定但计算成本越高（文档未明确提及，为时间序列预测常见设置）
            # 其他默认参数
            self.root_path = './datasets/'  # 数据集根目录路径，指定存储所有数据集文件的文件夹；取值范围：根据本地文件系统设置，如'./data/'、'/home/user/datasets/'等
            self.data_path = 'ETTh1.csv'  # 初始数据集文件名，会被data_parser根据self.data自动覆盖为对应数据集文件名；取值范围：与self.data匹配，如self.data='ETTm1'时覆盖为'ETTm1.csv' {insert\_element\_5\_}
            self.data_split = '0.7,0.1,0.2'  # 数据集划分比例（训练集:验证集:测试集），ILI和Traffic默认此比例，ETTh1等前四个数据集用时间顺序划分（如ETTh1为12/4/4个月）；取值范围：“train,val,test”格式字符串，常用{0.7,0.1,0.2}、{0.6,0.2,0.2} {insert\_element\_6\_}
            self.checkpoints = './checkpoints/'  # 模型检查点保存路径，存储训练中最优模型参数；取值范围：根据本地文件系统设置，如'./model_checkpoints/'等
            self.win_size = 2  # 分层编解码器（HED）中分段合并的窗口大小，控制相邻分段合并范围；取值范围：{2,4}，需与分段合并逻辑匹配（如窗口2表示合并2个相邻分段） {insert\_element\_7\_}
            self.factor = 10  # TSA层路由机制的路由器数量相关系数，默认与文档中最优路由器数量c=10匹配；取值范围：{5,10,15}，c≥5时长期预测精度稳定 {insert\_element\_8\_}
            self.data_dim = 7  # 数据集维度数量（D），ETTh1含7个电力变压器指标，会被data_parser根据实际数据集覆盖；取值范围：等于对应数据集维度数，如ETTh1/ETTm1=7、WTH=12、ECL=321、ILI=7、Traffic=862 {insert\_element\_9\_}
            self.d_model = 256  # 模型隐藏层维度（d_model），即DSW嵌入后向量维度，ETTh1等低维数据集用256；取值范围：低维数据集（D≤12）取{128,256}，高维数据集（D≥300）取{64,128} {insert\_element\_10\_}
            self.d_ff = 512  # 前馈网络（MLP）隐藏层维度，通常为d_model的2-4倍；取值范围：{256,512,1024}，如d_model=256时取512，d_model=64时取128 {insert\_element\_11\_}
            self.n_heads = 4  # 多头注意力（MSA）头数，ETTh1等低维数据集用4头；取值范围：低维数据集取{4,8}，高维数据集取{2,4}，需满足d_model能被头数整除 {insert\_element\_12\_}
            self.e_layers = 3  # 编码器层数（N），文档默认3层，通过分层结构捕捉不同尺度依赖；取值范围：{2,3,4}，层数越多可捕捉复杂依赖但计算成本增加 {insert\_element\_13\_}
            self.dropout = 0.2  # Dropout正则化概率，防止过拟合；取值范围：{0.1,0.2,0.3}，值过高丢失信息，过低易过拟合（深度学习常见设置）
            self.baseline = False  # 是否使用基线模型（如原始Transformer）对比，True表示运行基线模型，False表示运行Crossformer；取值范围：True/False {insert\_element\_14\_}
            self.num_workers = 0  # 数据加载线程数，0表示不使用多线程（适用于Windows或小数据量）；取值范围：{0,2,4,8}，根据CPU核心数调整
            self.batch_size = 32  # 训练批次大小，文档默认32，控制每次迭代样本数量；取值范围：{16,32,64}，GPU内存充足取大值，不足取小值 {insert\_element\_15\_}
            self.train_epochs = 15  # 总训练轮数，文档默认20轮；取值范围：{15,20,25}，结合早停机制使用 {insert\_element\_16\_}
            self.patience = 3  # 早停机制耐心值，验证损失连续3轮未降则提前停止；取值范围：{2,3,5}，值过小易未收敛，过大增加无效训练 {insert\_element\_17\_}
            self.lradj = 'type1'  # 学习率调整策略，'type1'通常为损失触发衰减；取值范围：'type1'（损失触发）、'type2'（固定步数）、'cos'（余弦退火）（常见设置）
            self.save_pred = False  # 是否保存预测结果用于后续分析（如可视化）；取值范围：True/False，True表示保存结果 {insert\_element\_18\_}
            self.use_gpu = True  # 是否使用GPU训练，文档实验均在GPU运行以加速；取值范围：True/False，True用GPU，False用CPU（仅调试或小数据量） {insert\_element\_19\_}
            self.gpu = 0  # 单GPU场景下指定设备编号，use_multi_gpu为False时生效；取值范围：GPU设备编号，如0、1等（根据实际GPU数量调整）
            self.use_multi_gpu = False  # 是否使用多GPU并行训练；取值范围：True/False，True启用多GPU（代码扩展功能）
            self.devices = '0,1'  # 多GPU场景下指定设备编号，use_multi_gpu为True时生效；取值范围：GPU编号组合字符串，如'0,1'、'1,2,3'等

    # 初始化参数
    args = Args()

    # 处理GPU设置
    args.use_gpu = True if torch.cuda.is_available() and args.use_gpu else False
    if args.use_gpu and args.use_multi_gpu:
        args.devices = args.devices.replace(' ', '')
        device_ids = args.devices.split(',')
        args.device_ids = [int(id_) for id_ in device_ids]
        args.gpu = args.device_ids[0]
        print(f"使用主GPU: {args.gpu}")

    # 数据解析器（自动设置数据集相关参数）
    data_parser = {
        'ETTh1': {'data': 'ETTh1.csv', 'data_dim': 7, 'split': [12 * 30 * 24, 4 * 30 * 24, 4 * 30 * 24]},
        'ETTm1': {'data': 'ETTm1.csv', 'data_dim': 7, 'split': [4 * 12 * 30 * 24, 4 * 4 * 30 * 24, 4 * 4 * 30 * 24]},
        'WTH': {'data': 'WTH.csv', 'data_dim': 12, 'split': [28 * 30 * 24, 10 * 30 * 24, 10 * 30 * 24]},
        'ECL': {'data': 'ECL.csv', 'data_dim': 321, 'split': [15 * 30 * 24, 3 * 30 * 24, 4 * 30 * 24]},
        'ILI': {'data': 'national_illness.csv', 'data_dim': 7, 'split': [0.7, 0.1, 0.2]},
        'Traffic': {'data': 'traffic.csv', 'data_dim': 862, 'split': [0.7, 0.1, 0.2]},
    }
    if args.data in data_parser.keys():
        data_info = data_parser[args.data]
        args.data_path = data_info['data']
        args.data_dim = data_info['data_dim']
        args.data_split = data_info['split']
    else:
        args.data_split = string_split(args.data_split)

    # 打印参数
    print('实验参数:')
    for arg in vars(args):
        print(f"{arg}: {getattr(args, arg)}")


    # 运行实验
    Exp = Exp_crossformer
    for ii in range(args.itr):
        setting = (f'Crossformer_{args.data}_il{args.in_len}_ol{args.out_len}_sl{args.seg_len}_'
                   f'win{args.win_size}_fa{args.factor}_dm{args.d_model}_nh{args.n_heads}_'
                   f'el{args.e_layers}_itr{ii}')

        exp = Exp(args)
        print(f'>>>>>>>开始训练: {setting}>>>>>>>>>>>>>>>>>>>>>>>>>>')
        exp.train(setting)

        print(f'>>>>>>>开始测试: {setting}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<')
        exp.test(setting, args.save_pred)


if __name__ == "__main__":
    print(666)
    print(torch.__version__)
    print(torch.version.cuda)
    print(torch.cuda.is_available())
    argDefin()